package cn.edu.xjtu.competiton.力扣.周赛.weekly291;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

/**
 * @author Hydrion-QLz
 * @date 2022-05-01 10:28
 * @description https://leetcode-cn.com/contest/weekly-contest-291/problems/k-divisible-elements-subarrays/
 * <p>
 * 给你一个整数数组 nums 和两个整数 k 和 p ，找出并返回满足要求的不同的子数组数，要求子数组中最多 k 个可被 p 整除的元素。
 * <p>
 * 如果满足下述条件之一，则认为数组 nums1 和 nums2 是 不同 数组：
 * <p>
 * 两数组长度 不同 ，或者
 * 存在 至少 一个下标 i 满足 nums1[i] != nums2[i] 。
 * 子数组 定义为：数组中的连续元素组成的一个 非空 序列。
 */
public class third {
    public static void main(String[] args) {
        third solution = new third();
        int[] arr = {2, 3, 3, 2, 2};
        System.out.println(solution.countDistinct(arr, 2, 2));
    }

    public int countDistinct(int[] nums, int k, int p) {
        int n = nums.length;
        int[] prefix = new int[n + 1];
        int num = 0;
        for (int i = 0; i < nums.length; i++) {
            prefix[i] = num;
            num += nums[i] % p == 0 ? 1 : 0;
        }
        prefix[n] = num;
        int ans = 0;
        Map<Integer, ArrayList<int[]>> map = new HashMap<>();
        for (int i = 0; i < prefix.length; i++) {
            for (int j = i + 1; j < prefix.length; j++) {
                if (prefix[j] - prefix[i] <= k && judge(map, i, j, nums)) {
                    ans++;
                    ArrayList<int[]> list = map.getOrDefault(j - i, new ArrayList<>());
                    list.add(new int[]{i, j});
                    map.put(j - i, list);
                }
            }
        }
        return ans;
    }

    private boolean judge(Map<Integer, ArrayList<int[]>> map, int i, int j, int[] nums) {
        ArrayList<int[]> list = map.get(j - i);
        if (list == null) {
            return true;
        }
        for (int[] point : list) {
            if (Arrays.equals(nums, i, j, nums, point[0], point[1])) {
                return false;
            }
        }
        return true;
    }
}
